import torch

from model.CNN import CNN
from util.webUtils import jsonToTensor


# 模型预测
def modelPredictServer(segmentData):
    # json转tensor
    segment_tensor=jsonToTensor(segmentData)

    model = CNN(segment_tensor.shape, 6)
    model.load_state_dict(torch.load('saved/trained_model.pth'))
    sleep_stages = ['wakefulness', 'light_sleep', 'deep_sleep', 'deep_sleep', 'REM']
    model.eval()
    with torch.no_grad():
        outputs = model(segment_tensor)
        _, predicted_class = torch.max(outputs, 1)
        predicted_stage = sleep_stages[predicted_class.item() - 1]
        print("Predicted sleep stage for the segment:", predicted_stage)
    return predicted_stage


